#!/usr/bin/python3
# -*- coding:utf-8 -*-
# __author__ = '__Jack__'

from typing import List
import time
import random

import requests
from fastapi.websockets import WebSocket
from fastapi import APIRouter, Depends, HTTPException, BackgroundTasks, Request
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from pydantic import HttpUrl
from sqlalchemy.orm import Session

from st.coronavirus import crud, schemas
from st.coronavirus.database import engine, Base, SessionLocal
from st.coronavirus.models import City, Data

application = APIRouter()

templates = Jinja2Templates(directory='../st/coronavirus/templates')

Base.metadata.create_all(bind=engine)


def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


@application.post("/create_city", response_model=schemas.ReadCity)
def create_city(city: schemas.CreateCity, db: Session = Depends(get_db)):
    db_city = crud.get_city_by_name(db, name=city.province)
    if db_city:
        raise HTTPException(status_code=400, detail="City already registered")
    return crud.create_city(db=db, city=city)


@application.get("/get_city/{city}", response_model=schemas.ReadCity)
def get_city(city: str, db: Session = Depends(get_db)):
    db_city = crud.get_city_by_name(db, name=city)
    if db_city is None:
        raise HTTPException(status_code=404, detail="City not found")
    return db_city


@application.get("/get_cities", response_model=List[schemas.ReadCity])
async def get_cities(skip: int = 0, limit: int = 100, db: Session = Depends(get_db)):
    cities = crud.get_cities(db, skip=skip, limit=limit)
    return cities


@application.post("/create_data", response_model=schemas.ReadData)
def create_data_for_city(city: str, data: schemas.CreateData, db: Session = Depends(get_db)):
    db_city = crud.get_city_by_name(db, name=city)
    data = crud.create_city_data(db=db, data=data, city_id=db_city.id)
    return data


@application.get("/get_data")
async def get_data(city: str = None, skip: int = 0, limit: int = 100, db: Session = Depends(get_db)):
    data = crud.get_data(db, city=city, skip=skip, limit=limit)
    return data


@application.get("/get_city_num", response_model=schemas.GetCityNumber)
async def get_city_num(db: Session = Depends(get_db)):
    data = crud.get_city_num(db)
    return schemas.GetCityNumber(number=data)


def bg_task(url: HttpUrl, db: Session):
    """这里注意一个坑，不要在后台任务的参数中db: Session = Depends(get_db)这样导入依赖"""

    city_data = requests.get(url=f"{url}?source=jhu&country_code=CN&timelines=false")

    if 200 == city_data.status_code:
        db.query(City).delete()  # 同步数据前先清空原有的数据
        for location in city_data.json()["locations"]:
            city = {
                "province": location["province"],
                "country": location["country"],
                "country_code": "CN",
                "country_population": location["country_population"]
            }
            crud.create_city(db=db, city=schemas.CreateCity(**city))

    coronavirus_data = requests.get(url=f"{url}?source=jhu&country_code=CN&timelines=true")

    if 200 == coronavirus_data.status_code:
        db.query(Data).delete()
        for city in coronavirus_data.json()["locations"]:
            db_city = crud.get_city_by_name(db=db, name=city["province"])
            for date, confirmed in city["timelines"]["confirmed"]["timeline"].items():
                data = {
                    "date": date.split("T")[0],  # 把'2020-12-31T00:00:00Z' 变成 ‘2020-12-31’
                    "confirmed": confirmed,
                    "deaths": city["timelines"]["deaths"]["timeline"][date],
                    "recovered": 0  # 每个城市每天有多少人痊愈，这种数据没有
                }
                # 这个city_id是city表中的主键ID，不是coronavirus_data数据里的ID
                crud.create_city_data(db=db, data=schemas.CreateData(**data), city_id=db_city.id)


@application.get("/sync_coronavirus_data/jhu")
def sync_coronavirus_data(background_tasks: BackgroundTasks, db: Session = Depends(get_db)):
    """从Johns Hopkins University同步COVID-19数据"""
    background_tasks.add_task(bg_task, "https://coronavirus-tracker-api.herokuapp.com/v2/locations", db)
    return {"message": "正在后台同步数据..."}


@application.get("/")
def coronavirus(request: Request, city: str = None, skip: int = 0, limit: int = 100, db: Session = Depends(get_db)):
    data = crud.get_data(db, city=city, skip=skip, limit=limit)
    return templates.TemplateResponse("home.html", {
        "request": request,
        "data": data,
        "sync_data_url": "/coronavirus/sync_coronavirus_data/jhu"
    })


def bg_task_two(count: int, db: Session):
    city_list = []
    for i in range(count):
        city = {
            "province": str(i) + '_province_' + str(i + random.randint(1, 5055050)),
            "country": str(i) + '_country_' + str(i + random.randint(1, 5055050)),
            "country_code": "CN",
            "country_population": random.randint(1, 5055050)
        }
        city_list.append(schemas.CreateCity(**city))
    crud.create_batch_city(db=db, citys=city_list)


@application.get("/insert_data")
async def insert_data(background_tasks: BackgroundTasks, count: int = 100, db: Session = Depends(get_db)):
    background_tasks.add_task(bg_task_two, count, db)
    return {"message": "正在后台同步数据..."}


async def bg_task_three(count: int, db: Session, websocket: WebSocket):
    print(f"异步任务开始")
    city_list = []
    for i in range(count):
        city = {
            "province": str(i) + '_province_' + str(i + random.randint(1, 5055050)),
            "country": str(i) + '_country_' + str(i + random.randint(1, 5055050)),
            "country_code": "CN",
            "country_population": random.randint(1, 5055050)
        }

        city_list.append(schemas.CreateCity(**city))
        if i % 5 == 0:
            print(f"数据已生成： {i}")
            await websocket.send_text(f"已插入数据: {i} 条数据")
            time.sleep(2)
    crud.create_batch_city(db=db, citys=city_list)


@application.websocket("/ws")
async def ws(websocket: WebSocket, db: Session = Depends(get_db)):
    await websocket.accept()
    while True:
        # websocket.receive_bytes()
        # websocket.receive_json()
        data = await websocket.receive_text()
        print(data)
        if data.isdigit():
            print(f"插入数据开始~~~~~~~~~~: {data}")
            await websocket.send_text(f"收到来自客户端的请求插入: {data} 条数据")
            await bg_task_three(int(data), db, websocket)
